from utils import date_util
from data.read_data import get_data_by_pickle, get_data_for_model
from data.time_data import TimeData
from utils import date_util
from strategy import strategy_base

class BackTester(object):
    def __init__(self, trade_strategy: strategy_base, start_date, end_date):
        self.trade_strategy = trade_strategy
        if hasattr(self.trade_strategy, 'recent_day'):
            recent_day = trade_strategy.recent_day
        else:
            recent_day = 0
        self.time_data = TimeData(date_util.add_some_day(start_date, -recent_day), date_util.add_one_day(end_date),
                                  get_data_by_pickle(), get_data_for_model())
        self.time_data.start_date = start_date
        self.time_data.end_date = end_date
        # 交易盈亏结果序列
        self.money_array = {}
        self.capital_array = {}
        self.profit_array={}
        self.stock_array=self.trade_strategy.stock_array
    def execute_trade(self):
        start_date, end_date = self.time_data.start_date, self.time_data.end_date
        date = start_date
        while date != end_date:
            try:
                # hasattr: 用来查询对象有没有实现某个方法
                if hasattr(self.trade_strategy, 'buy_strategy'):
                    # 买入策略执行
                    self.trade_strategy.buy_strategy(time_data=self.time_data, now_date=date, start_date=start_date)

                if hasattr(self.trade_strategy, 'sell_strategy'):
                    # 卖出策略执行
                    self.trade_strategy.sell_strategy(time_data=self.time_data, now_date=date, start_date=start_date,
                                                      end_flag=False)

                self.money_array[date]=self.trade_strategy.money_now

                equity=0
                index = date_util.date_str_diff(date, start_date)
                if date in self.stock_array:
                    for stock in self.stock_array[date]:
                        equity+=self.time_data.stock_prices[stock[0]][index]*stock[-1]
                #print("equity:",equity)
                self.capital_array[date]=self.money_array[date]+equity

                yesterday=date_util.add_some_day(date,-1)
                self.profit_array[date]=self.capital_array[date]-self.capital_array.get(yesterday,1000)

                date = date_util.add_one_day(date)
            except IndexError:
                date = date_util.add_one_day(date)
                continue
        if hasattr(self.trade_strategy, 'sell_strategy'):
            # 卖出策略执行
            self.trade_strategy.sell_strategy(time_data=self.time_data, now_date=date, start_date=start_date,
                                              end_flag=True)
    
    def calculate_profit(self):
        return self.money_array[-1]-self.money_array[0]

    def get_stock_array(self): #返回每天的持仓情况
        #return type:dict
        return self.stock_array

    def get_cash_array(self): #返回每天剩余的现金
        #return type:dict
        return self.money_array
    
    def get_capital_array(self): #返回每天现金和股票的总金额
        #return type:dict
        return self.capital_array
    
    def get_profit_array(self): #返回每天的收益
        #return type:dict
        return self.profit_array
